package com.wgz.flag.websocket;

import jakarta.websocket.Session;
import java.io.IOException;
import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

@Slf4j
@Component
public class WebSocketSessionManager {
    // 用户ID -> WebSocket会话
    private static final Map<Long, Session> USER_SESSIONS = new ConcurrentHashMap<>();
    // 群ID -> 用户ID集合
    private static final Map<Long, Set<Long>> GROUP_USERS = new ConcurrentHashMap<>();
    // Session -> 用户ID
    private static final Map<Session, Long> SESSION_USER_MAP = new ConcurrentHashMap<>();

    public static void addSession(Long userId, Session session) {
        USER_SESSIONS.put(userId, session);
        SESSION_USER_MAP.put(session, userId);
        log.info("Added session for user {}. Total active sessions: {}", userId, USER_SESSIONS.size());
    }

    public static void removeSession(Long userId) {
        Session session = USER_SESSIONS.remove(userId);
        if (session != null) {
            SESSION_USER_MAP.remove(session);
        }
        log.info("用户[{}]断开连接", userId);
    }

    public static void joinGroup(Long groupId, Long userId) {
        GROUP_USERS.computeIfAbsent(groupId, k -> ConcurrentHashMap.newKeySet()).add(userId);
    }

    public static void leaveGroup(Long groupId, Long userId) {
        Set<Long> users = GROUP_USERS.get(groupId);
        if (users != null) {
            users.remove(userId);
        }
    }

    public static void sendToUser(Long userId, String message) {
        Session session = USER_SESSIONS.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                log.error("发送消息给用户[{}]失败", userId, e);
            }
        }
    }

    public static void sendToGroup(Long groupId, String message, Long excludeUserId) {
        Set<Long> users = GROUP_USERS.get(groupId);
        if (users != null) {
            users.stream()
                    .filter(userId -> !userId.equals(excludeUserId))
                    .forEach(userId -> sendToUser(userId, message));
        }
    }

    public static boolean isOnline(Long userId) {
        Session session = USER_SESSIONS.get(userId);
        return session != null && session.isOpen();
    }

    public static Set<Long> getGroupUsers(Long groupId) {
        return GROUP_USERS.getOrDefault(groupId, ConcurrentHashMap.newKeySet());
    }

    public static Collection<Session> getAllSessions() {
        return USER_SESSIONS.values();
    }

    public static Session getSession(Long userId) {
        Session session = USER_SESSIONS.get(userId);
        if (session != null) {
            // 检查会话是否真的可用
            if (!session.isOpen()) {
                log.warn("Found closed session for user {}, removing it", userId);
                removeSession(userId);
                return null;
            }
            return session;
        }
        return null;
    }

    public static Long getUserIdBySession(Session session) {
        for (Map.Entry<Long, Session> entry : USER_SESSIONS.entrySet()) {
            if (entry.getValue().equals(session)) {
                return entry.getKey();
            }
        }
        return null;
    }
}